import logging
import os
import cv2
import numpy as np
from PIL import Image
import io
import base64
from tools._dectimer import timer
from tools.log_util import INFO, ERROR, WARNING

try:
    # 尝试导入 pytesseract，但不是必须的
    import pytesseract
    TESSERACT_AVAILABLE = True
except ImportError:
    TESSERACT_AVAILABLE = False

# logger = logging.getLogger(__name__)

@timer()
def preprocess_image(image_path):
    """
    预处理验证码图片以提高识别率
    
    Args:
        image_path: 图片路径
        
    Returns:
        处理后的图片路径
    """
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")
    
    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 二值化处理
    _, binary = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)
    
    # 去噪
    kernel = np.ones((2, 2), np.uint8)
    opening = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
    
    # 膨胀，使文字更清晰
    dilation = cv2.dilate(opening, kernel, iterations=1)
    
    # 保存处理后的图片
    processed_path = f"{os.path.splitext(image_path)[0]}_processed.png"
    cv2.imwrite(processed_path, dilation)
    
    return processed_path

@timer()
def recognize_captcha_with_tesseract(image_path, preprocess=True):
    """
    使用 Tesseract OCR 识别验证码图片中的数字
    
    Args:
        image_path: 图片路径
        preprocess: 是否进行预处理
        
    Returns:
        识别出的文本
    """
    if not TESSERACT_AVAILABLE:
        WARNING.logger.warning("Tesseract OCR 未安装，无法使用此方法识别验证码")
        return None
        
    try:
        if preprocess:
            processed_path = preprocess_image(image_path)
            img = Image.open(processed_path)
        else:
            img = Image.open(image_path)
            
        # 设置tesseract配置，只识别数字
        custom_config = r'--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789'
        
        # 识别图片中的文字
        text = pytesseract.image_to_string(img, config=custom_config)
        
        # 清理结果，只保留数字
        result = ''.join(c for c in text if c.isdigit())
        
        INFO.logger.info(f"Tesseract OCR 验证码识别结果: {result}")
        return result
    except Exception as e:
        ERROR.logger.error(f"Tesseract OCR 验证码识别失败: {e}")
        return None

@timer()
def recognize_captcha_with_opencv(image_path, preprocess=True):
    """
    使用 OpenCV 进行简单的数字验证码识别
    这种方法对于简单的数字验证码有一定效果，但对复杂验证码效果有限
    
    Args:
        image_path: 图片路径
        preprocess: 是否进行预处理
        
    Returns:
        识别出的文本
    """
    try:
        if preprocess:
            processed_path = preprocess_image(image_path)
            img = cv2.imread(processed_path)
        else:
            img = cv2.imread(image_path)
            
        # 转为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # 二值化
        _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)
        
        # 查找轮廓
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 根据轮廓分割字符
        digit_images = []
        for contour in sorted(contours, key=lambda c: cv2.boundingRect(c)[0]):
            x, y, w, h = cv2.boundingRect(contour)
            # 过滤掉太小的轮廓
            if w > 5 and h > 10:
                digit_images.append((x, thresh[y:y+h, x:x+w]))
        
        # 对分割出的字符进行简单识别
        # 这里只是一个非常简单的示例，实际应用中可能需要更复杂的识别算法
        result = ""
        for x, digit_img in sorted(digit_images, key=lambda d: d[0]):
            # 根据像素密度等特征进行简单判断
            # 这里只是一个非常简化的示例
            pixel_sum = np.sum(digit_img)
            height, width = digit_img.shape
            density = pixel_sum / (height * width * 255)
            
            # 打印调试信息
            INFO.logger.debug(f"字符宽度: {width}, 高度: {height}, 像素密度: {density:.2f}")
            
            # 这里需要根据实际验证码特征调整判断逻辑
            # 下面的逻辑只是示例，实际应用需要根据验证码特点调整
            if 0.3 <= density <= 0.35:
                result += "1"
            elif 0.35 < density <= 0.45:
                result += "7"
            elif 0.45 < density <= 0.55:
                result += "4"
            elif 0.55 < density <= 0.65:
                result += "0"
            elif 0.65 < density <= 0.75:
                result += "8"
            else:
                result += "?"
        
        INFO.logger.info(f"OpenCV 验证码识别结果: {result}")
        return result
    except Exception as e:
        ERROR.logger.error(f"OpenCV 验证码识别失败: {e}")
        return None

@timer()
def recognize_captcha_manually(image_path):
    """
    手动识别验证码（显示图片并等待用户输入）
    
    Args:
        image_path: 图片路径
        
    Returns:
        用户输入的验证码
    """
    try:
        # 显示图片
        img = cv2.imread(image_path)
        cv2.imshow('Captcha', img)
        print("请查看验证码图片并在控制台输入验证码（按任意键关闭图片窗口）")
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        
        # 获取用户输入
        captcha = input("请输入验证码: ")
        result = ''.join(c for c in captcha if c.isdigit())
        
        INFO.logger.info(f"手动验证码识别结果: {result}")
        return result
    except Exception as e:
        ERROR.logger.error(f"手动验证码识别失败: {e}")
        return None

@timer()
def recognize_captcha(image_path, preprocess=True, method='auto'):
    """
    识别验证码图片中的数字
    
    Args:
        image_path: 图片路径
        preprocess: 是否进行预处理
        method: 识别方法，可选值：'tesseract', 'opencv', 'manual', 'auto'
        
    Returns:
        识别出的文本
    """
    INFO.logger.info(f"开始识别验证码图片: {image_path}，方法: {method}")
    
    result = None
    
    if method == 'tesseract' or method == 'auto':
        if TESSERACT_AVAILABLE:
            result = recognize_captcha_with_tesseract(image_path, preprocess)
            if result and method == 'tesseract':
                return result
    
    if method == 'opencv' or (method == 'auto' and not result):
        result = recognize_captcha_with_opencv(image_path, preprocess)
        if result and method == 'opencv':
            return result
    
    if method == 'manual' or (method == 'auto' and not result):
        result = recognize_captcha_manually(image_path)
    
    return result

if __name__ == '__main__':
    # 测试验证码识别
    captcha_path = 'tools/captcha.png'
    result = recognize_captcha(captcha_path)
    print(f"识别结果: {result}") 